"""
NS-SEM Solver Validator Implementation

Handles validation of NS-SEM solver results and configurations.
"""

import os
import re
import math
from pathlib import Path
from typing import Dict, List, Any, Optional, Tuple
import logging
import xml.etree.ElementTree as ET

from gcr_solver_manager.extensions.base_validator import BaseValidator, ValidationResult

logger = logging.getLogger(__name__)


class NSsemValidator(BaseValidator):
    """
    Validator for NS-SEM (Navier-Stokes Spectral Element Method) solver.
    
    Handles validation of solver results, numerical accuracy, conservation
    laws, and configuration correctness for NS-SEM simulations.
    """
    
    def __init__(self, extension_name: str):
        super().__init__(extension_name)
        self.default_tolerance = 1e-6
        self.conservation_tolerance = 1e-8
    
    def can_validate(self, output_path: str) -> bool:
        """
        Check if this validator can validate the given output.
        
        Args:
            output_path: Path to output file or directory
            
        Returns:
            True if validator can handle this output
        """
        output_path_obj = Path(output_path)
        
        # Check for NS-SEM specific output patterns
        if output_path_obj.is_file():
            name = output_path_obj.name.lower()
            ns_sem_patterns = [
                "ns_sem",
                "tensor_mat", 
                "multi_cavity",
                "spectral_element",
                "navier_stokes"
            ]
            
            for pattern in ns_sem_patterns:
                if pattern in name:
                    return True
        
        elif output_path_obj.is_dir():
            # Check for NS-SEM output files in directory
            ns_sem_files = [
                "ns_sem*.log",
                "tensor_*.out",
                "cavity_*.dat",
                "residuals_*.txt"
            ]
            
            for pattern in ns_sem_files:
                if list(output_path_obj.glob(pattern)):
                    return True
        
        return False
    
    def validate_convergence(self, output_data: str, config: Dict[str, Any]) -> ValidationResult:
        """
        Validate solver convergence.
        
        Args:
            output_data: Solver output text
            config: Validation configuration
            
        Returns:
            Validation result for convergence
        """
        logger.debug("Validating NS-SEM solver convergence")
        
        tolerance = config.get("tolerance", self.default_tolerance)
        max_iterations = config.get("max_iterations", 1000)
        
        # Parse convergence information
        convergence_info = self._parse_convergence_info(output_data)
        
        if not convergence_info:
            return ValidationResult(
                success=False,
                message="No convergence information found in output",
                details={"error": "convergence_info_missing"},
                metrics={}
            )
        
        # Check if converged
        converged = convergence_info.get("converged", False)
        iterations = convergence_info.get("iterations", 0)
        final_residual = convergence_info.get("final_residual", float('inf'))
        
        success = True
        issues = []
        
        # Check convergence status
        if not converged:
            success = False
            issues.append(f"Solver did not converge after {iterations} iterations")
        
        # Check final residual
        if final_residual > tolerance:
            success = False
            issues.append(f"Final residual {final_residual:.2e} exceeds tolerance {tolerance:.2e}")
        
        # Check iteration count
        if iterations >= max_iterations:
            success = False
            issues.append(f"Reached maximum iterations ({max_iterations})")
        
        message = "Convergence validated successfully"
        if issues:
            message = f"Convergence validation failed: {'; '.join(issues)}"
        
        metrics = {
            "converged": converged,
            "iterations": iterations,
            "final_residual": final_residual,
            "tolerance": tolerance,
            "convergence_rate": convergence_info.get("convergence_rate", 0.0)
        }
        
        return ValidationResult(
            success=success,
            message=message,
            details=convergence_info,
            metrics=metrics
        )
    
    def validate_conservation(self, output_data: str, config: Dict[str, Any]) -> ValidationResult:
        """
        Validate conservation laws (mass, momentum, energy).
        
        Args:
            output_data: Solver output text
            config: Validation configuration
            
        Returns:
            Validation result for conservation laws
        """
        logger.debug("Validating NS-SEM conservation laws")
        
        tolerance = config.get("conservation_tolerance", self.conservation_tolerance)
        
        # Parse conservation data
        conservation_data = self._parse_conservation_data(output_data)
        
        if not conservation_data:
            return ValidationResult(
                success=False,
                message="No conservation data found in output",
                details={"error": "conservation_data_missing"},
                metrics={}
            )
        
        success = True
        issues = []
        metrics = {}
        
        # Check mass conservation
        if "mass_conservation" in conservation_data:
            mass_error = abs(conservation_data["mass_conservation"])
            metrics["mass_conservation_error"] = mass_error
            
            if mass_error > tolerance:
                success = False
                issues.append(f"Mass conservation error {mass_error:.2e} exceeds tolerance {tolerance:.2e}")
        
        # Check momentum conservation
        momentum_components = ["momentum_x", "momentum_y", "momentum_z"]
        for component in momentum_components:
            if component in conservation_data:
                momentum_error = abs(conservation_data[component])
                metrics[f"{component}_error"] = momentum_error
                
                if momentum_error > tolerance:
                    success = False
                    issues.append(f"{component} conservation error {momentum_error:.2e} exceeds tolerance")
        
        # Check energy conservation
        if "energy_conservation" in conservation_data:
            energy_error = abs(conservation_data["energy_conservation"])
            metrics["energy_conservation_error"] = energy_error
            
            if energy_error > tolerance:
                success = False
                issues.append(f"Energy conservation error {energy_error:.2e} exceeds tolerance {tolerance:.2e}")
        
        message = "Conservation laws validated successfully"
        if issues:
            message = f"Conservation validation failed: {'; '.join(issues)}"
        
        return ValidationResult(
            success=success,
            message=message,
            details=conservation_data,
            metrics=metrics
        )
    
    def validate_numerical_accuracy(self, output_data: str, reference_data: Optional[str], config: Dict[str, Any]) -> ValidationResult:
        """
        Validate numerical accuracy against reference solution.
        
        Args:
            output_data: Solver output text
            reference_data: Reference solution data
            config: Validation configuration
            
        Returns:
            Validation result for numerical accuracy
        """
        logger.debug("Validating NS-SEM numerical accuracy")
        
        # Parse numerical data
        numerical_data = self._parse_numerical_data(output_data)
        
        if not numerical_data:
            return ValidationResult(
                success=False,
                message="No numerical data found in output",
                details={"error": "numerical_data_missing"},
                metrics={}
            )
        
        metrics = numerical_data.copy()
        success = True
        issues = []
        
        # Check solution norms
        if "solution_norm" in numerical_data:
            solution_norm = numerical_data["solution_norm"]
            metrics["solution_l2_norm"] = solution_norm
            
            # Check for NaN or Inf values
            if not math.isfinite(solution_norm):
                success = False
                issues.append(f"Solution norm is not finite: {solution_norm}")
            elif solution_norm == 0.0:
                success = False
                issues.append("Solution norm is zero (trivial solution)")
        
        # Validate against reference if available
        if reference_data:
            reference_metrics = self._parse_numerical_data(reference_data)
            if reference_metrics:
                accuracy_result = self._compare_with_reference(numerical_data, reference_metrics, config)
                success = success and accuracy_result.success
                if not accuracy_result.success:
                    issues.append(accuracy_result.message)
                metrics.update(accuracy_result.metrics)
        
        # Check spectral accuracy indicators
        if "spectral_decay" in numerical_data:
            spectral_decay = numerical_data["spectral_decay"]
            if spectral_decay < config.get("min_spectral_decay", 1e-12):
                issues.append(f"Insufficient spectral decay: {spectral_decay:.2e}")
        
        message = "Numerical accuracy validated successfully"
        if issues:
            message = f"Numerical accuracy validation failed: {'; '.join(issues)}"
        
        return ValidationResult(
            success=success,
            message=message,
            details=numerical_data,
            metrics=metrics
        )
    
    def validate_configuration(self, config_path: str) -> ValidationResult:
        """
        Validate solver configuration file.
        
        Args:
            config_path: Path to configuration file
            
        Returns:
            Validation result for configuration
        """
        logger.debug(f"Validating NS-SEM configuration: {config_path}")
        
        config_path_obj = Path(config_path)
        
        if not config_path_obj.exists():
            return ValidationResult(
                success=False,
                message=f"Configuration file not found: {config_path}",
                details={"error": "config_file_missing"},
                metrics={}
            )
        
        try:
            # Parse configuration
            if config_path_obj.suffix.lower() == '.xml':
                config_data = self._parse_xml_config(config_path_obj)
            else:
                config_data = self._parse_text_config(config_path_obj)
            
            # Validate configuration parameters
            validation_issues = []
            
            # Check required parameters
            required_params = [
                "problem_type",
                "spectral_order", 
                "grid_nx",
                "solver_type",
                "max_iterations"
            ]
            
            for param in required_params:
                if param not in config_data:
                    validation_issues.append(f"Missing required parameter: {param}")
            
            # Validate parameter values
            if "spectral_order" in config_data:
                order = config_data["spectral_order"]
                if not isinstance(order, int) or order < 2 or order > 20:
                    validation_issues.append(f"Invalid spectral order: {order} (must be 2-20)")
            
            if "reynolds_number" in config_data:
                re_num = config_data["reynolds_number"]
                if not isinstance(re_num, (int, float)) or re_num <= 0:
                    validation_issues.append(f"Invalid Reynolds number: {re_num} (must be positive)")
            
            if "time_step" in config_data:
                dt = config_data["time_step"]
                if not isinstance(dt, (int, float)) or dt <= 0:
                    validation_issues.append(f"Invalid time step: {dt} (must be positive)")
            
            # Check grid dimensions
            grid_params = ["grid_nx", "grid_ny", "grid_nz"]
            for param in grid_params:
                if param in config_data:
                    value = config_data[param]
                    if not isinstance(value, int) or value < 1:
                        validation_issues.append(f"Invalid {param}: {value} (must be positive integer)")
            
            success = len(validation_issues) == 0
            message = "Configuration validated successfully"
            if validation_issues:
                message = f"Configuration validation failed: {'; '.join(validation_issues)}"
            
            return ValidationResult(
                success=success,
                message=message,
                details=config_data,
                metrics={"validation_issues_count": len(validation_issues)}
            )
            
        except Exception as e:
            logger.error(f"Error validating configuration: {e}")
            return ValidationResult(
                success=False,
                message=f"Configuration validation error: {str(e)}",
                details={"error": str(e)},
                metrics={}
            )
    
    def validate_environment(self) -> ValidationResult:
        """
        Validate the execution environment.
        
        Returns:
            Validation result for environment
        """
        logger.debug("Validating NS-SEM execution environment")
        
        issues = []
        metrics = {}
        
        # Check PETSc
        petsc_dir = os.getenv("PETSC_DIR")
        if petsc_dir and Path(petsc_dir).exists():
            metrics["petsc_available"] = True
            metrics["petsc_dir"] = petsc_dir
        else:
            issues.append("PETSc not available (PETSC_DIR not set or invalid)")
            metrics["petsc_available"] = False
        
        # Check Kokkos
        kokkos_root = os.getenv("Kokkos_ROOT")
        kokkos_available = False
        if kokkos_root and Path(kokkos_root).exists():
            kokkos_available = True
        elif Path("/usr/local/include/Kokkos_Core.hpp").exists():
            kokkos_available = True
        
        metrics["kokkos_available"] = kokkos_available
        if not kokkos_available:
            issues.append("Kokkos not available (may impact GPU performance)")
        
        # Check MPI
        import shutil
        mpi_available = shutil.which("mpirun") is not None
        metrics["mpi_available"] = mpi_available
        if not mpi_available:
            issues.append("MPI not available (parallel execution not possible)")
        
        # Check available memory
        try:
            import psutil
            available_memory = psutil.virtual_memory().available / (1024**3)  # GB
            metrics["available_memory_gb"] = available_memory
            if available_memory < 2.0:
                issues.append(f"Low available memory: {available_memory:.1f} GB")
        except ImportError:
            pass
        
        success = len(issues) == 0
        message = "Environment validation successful"
        if issues:
            message = f"Environment validation issues: {'; '.join(issues)}"
        
        return ValidationResult(
            success=success,
            message=message,
            details={"issues": issues},
            metrics=metrics
        )
    
    def _parse_convergence_info(self, output_data: str) -> Dict[str, Any]:
        """Parse convergence information from output."""
        convergence_info = {}
        
        # Look for convergence patterns
        lines = output_data.split('\n')
        
        for line in lines:
            line = line.strip()
            
            # Check for convergence status
            if "converged" in line.lower():
                convergence_info["converged"] = "failed" not in line.lower()
            
            # Extract iteration count
            iteration_match = re.search(r'(?:iteration|iter)\s*[:,]?\s*(\d+)', line.lower())
            if iteration_match:
                convergence_info["iterations"] = int(iteration_match.group(1))
            
            # Extract final residual
            residual_match = re.search(r'(?:residual|res)\s*[:=]?\s*([\d.]+e[+-]?\d+)', line.lower())
            if residual_match:
                convergence_info["final_residual"] = float(residual_match.group(1))
            
            # Extract convergence rate
            rate_match = re.search(r'(?:convergence\s*rate|rate)\s*[:=]?\s*([\d.]+)', line.lower())
            if rate_match:
                convergence_info["convergence_rate"] = float(rate_match.group(1))
        
        return convergence_info
    
    def _parse_conservation_data(self, output_data: str) -> Dict[str, Any]:
        """Parse conservation law data from output."""
        conservation_data = {}
        
        lines = output_data.split('\n')
        
        for line in lines:
            line = line.strip().lower()
            
            # Mass conservation
            mass_match = re.search(r'mass\s*(?:conservation|error)\s*[:=]?\s*([\d.e+-]+)', line)
            if mass_match:
                conservation_data["mass_conservation"] = float(mass_match.group(1))
            
            # Momentum conservation
            momentum_x_match = re.search(r'momentum[_\s]*x\s*(?:conservation|error)\s*[:=]?\s*([\d.e+-]+)', line)
            if momentum_x_match:
                conservation_data["momentum_x"] = float(momentum_x_match.group(1))
            
            momentum_y_match = re.search(r'momentum[_\s]*y\s*(?:conservation|error)\s*[:=]?\s*([\d.e+-]+)', line)
            if momentum_y_match:
                conservation_data["momentum_y"] = float(momentum_y_match.group(1))
            
            momentum_z_match = re.search(r'momentum[_\s]*z\s*(?:conservation|error)\s*[:=]?\s*([\d.e+-]+)', line)
            if momentum_z_match:
                conservation_data["momentum_z"] = float(momentum_z_match.group(1))
            
            # Energy conservation
            energy_match = re.search(r'energy\s*(?:conservation|error)\s*[:=]?\s*([\d.e+-]+)', line)
            if energy_match:
                conservation_data["energy_conservation"] = float(energy_match.group(1))
        
        return conservation_data
    
    def _parse_numerical_data(self, output_data: str) -> Dict[str, Any]:
        """Parse numerical data from output."""
        numerical_data = {}
        
        lines = output_data.split('\n')
        
        for line in lines:
            line = line.strip().lower()
            
            # Solution norm
            norm_match = re.search(r'(?:solution|l2)\s*norm\s*[:=]?\s*([\d.e+-]+)', line)
            if norm_match:
                numerical_data["solution_norm"] = float(norm_match.group(1))
            
            # Spectral decay
            spectral_match = re.search(r'spectral\s*decay\s*[:=]?\s*([\d.e+-]+)', line)
            if spectral_match:
                numerical_data["spectral_decay"] = float(spectral_match.group(1))
            
            # Maximum velocity
            vel_match = re.search(r'max\s*velocity\s*[:=]?\s*([\d.e+-]+)', line)
            if vel_match:
                numerical_data["max_velocity"] = float(vel_match.group(1))
            
            # Pressure information
            pressure_match = re.search(r'(?:pressure|p)\s*(?:max|norm)\s*[:=]?\s*([\d.e+-]+)', line)
            if pressure_match:
                numerical_data["pressure_norm"] = float(pressure_match.group(1))
        
        return numerical_data
    
    def _parse_xml_config(self, config_path: Path) -> Dict[str, Any]:
        """Parse XML configuration file."""
        config_data = {}
        
        tree = ET.parse(config_path)
        root = tree.getroot()
        
        # Extract parameters
        for param in root.findall(".//Parameter"):
            name = param.get("name")
            value = param.get("value") 
            param_type = param.get("type", "string")
            
            if name and value:
                # Convert value based on type
                if param_type == "int":
                    config_data[name] = int(value)
                elif param_type == "double" or param_type == "number":
                    config_data[name] = float(value)
                elif param_type == "bool":
                    config_data[name] = value.lower() in ["true", "1", "yes"]
                else:
                    config_data[name] = value
        
        return config_data
    
    def _parse_text_config(self, config_path: Path) -> Dict[str, Any]:
        """Parse text-based configuration file."""
        config_data = {}
        
        content = config_path.read_text()
        
        for line in content.split('\n'):
            line = line.strip()
            if line and not line.startswith('#') and '=' in line:
                key, value = line.split('=', 1)
                key = key.strip()
                value = value.strip()
                
                # Try to convert to appropriate type
                try:
                    if '.' in value:
                        config_data[key] = float(value)
                    else:
                        config_data[key] = int(value)
                except ValueError:
                    if value.lower() in ["true", "false"]:
                        config_data[key] = value.lower() == "true"
                    else:
                        config_data[key] = value
        
        return config_data
    
    def _compare_with_reference(self, current: Dict[str, Any], reference: Dict[str, Any], config: Dict[str, Any]) -> ValidationResult:
        """Compare current results with reference."""
        tolerance = config.get("reference_tolerance", 1e-3)
        
        success = True
        issues = []
        metrics = {}
        
        common_keys = set(current.keys()) & set(reference.keys())
        
        for key in common_keys:
            if isinstance(current[key], (int, float)) and isinstance(reference[key], (int, float)):
                current_val = float(current[key])
                ref_val = float(reference[key])
                
                if ref_val != 0:
                    relative_error = abs(current_val - ref_val) / abs(ref_val)
                else:
                    relative_error = abs(current_val - ref_val)
                
                metrics[f"{key}_relative_error"] = relative_error
                
                if relative_error > tolerance:
                    success = False
                    issues.append(f"{key}: relative error {relative_error:.2e} > tolerance {tolerance:.2e}")
        
        message = "Reference comparison successful" if success else f"Reference comparison failed: {'; '.join(issues)}"
        
        return ValidationResult(
            success=success,
            message=message,
            details={"compared_parameters": list(common_keys)},
            metrics=metrics
        )